#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import requests
import json
import time
import subprocess
import sys
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor
from urllib.parse import quote

# 配置信息
MODEL_NAME = "Qwen/Qwen3-0.6B-Base"
SAVE_PATH = "/home/zkh/lzt/damoxing/boshuyuce/xindaoyuce/Qwen3-0.6Bbase"
MAX_PARALLEL_TASKS = 6  # 最大并行任务数
TIMEOUT = 120  # 请求超时时间(秒)
RETRY_TIMES = 5  # 失败重试次数
CONTROL_FILE = os.path.join(SAVE_PATH, ".download_status.json")  # 下载状态记录文件

# 国内HuggingFace镜像站点
MIRRORS = [
    "https://hf-mirror.com",  # 首选镜像
    "https://mirror.ghproxy.com/https://huggingface.co",  # GitHub代理镜像
    "https://huggingface.co"  # 官方地址
]

# aria2配置 - 断点续传是关键配置
ARIA2_OPTS = {
    'max-connection-per-server': '16',  # 每个服务器连接数
    'split': '16',                      # 单文件分片数
    'min-split-size': '1M',             # 最小分片大小
    'max-concurrent-downloads': '5',     # 并发下载数
    'continue': 'true',                 # 断点续传（关键）
    'max-tries': str(RETRY_TIMES),      # 重试次数
    'retry-wait': '3',                  # 重试等待时间(秒)
    'connect-timeout': str(TIMEOUT),    # 连接超时
    'timeout': str(TIMEOUT),            # 整体超时
    'auto-file-renaming': 'false',      # 不重命名文件
    'allow-overwrite': 'true',          # 允许覆盖
    'console-log-level': 'notice',      # 日志级别
    'summary-interval': '15',           # 状态显示间隔(秒)
    'file-allocation': 'falloc',        # 文件分配方式，用于提高断点续传效率
    'disk-cache': '64M',                # 磁盘缓存大小，提高下载效率
}

def ensure_dir(directory):
    """确保目录存在，如果不存在则创建它"""
    if not os.path.exists(directory):
        os.makedirs(directory)
        print(f"创建目录: {directory}")

def load_download_status():
    """加载已下载文件状态"""
    if os.path.exists(CONTROL_FILE):
        try:
            with open(CONTROL_FILE, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception as e:
            print(f"无法加载下载状态文件: {e}")
    return {"downloaded": [], "failed": [], "last_update": ""}

def save_download_status(status):
    """保存下载状态"""
    status["last_update"] = time.strftime("%Y-%m-%d %H:%M:%S")
    try:
        with open(CONTROL_FILE, 'w', encoding='utf-8') as f:
            json.dump(status, f, ensure_ascii=False, indent=2)
    except Exception as e:
        print(f"保存下载状态失败: {e}")

def get_file_list_from_huggingface():
    """从HuggingFace获取模型文件列表，尝试多个镜像源"""
    file_list = []
    
    for mirror in MIRRORS:
        try:
            api_url = f"{mirror}/api/models/{MODEL_NAME}/tree/main"
            print(f"尝试从 {mirror} 获取文件列表...")
            
            headers = {
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
            }
            
            response = requests.get(api_url, headers=headers, timeout=TIMEOUT)
            response.raise_for_status()
            
            # HuggingFace API返回的是一个文件对象列表
            files = response.json()
            
            if not files or not isinstance(files, list):
                print(f"从 {mirror} 获取的文件列表格式不符合预期")
                continue
                
            print(f"成功从 {mirror} 获取文件列表，共 {len(files)} 个文件")
            
            # 处理文件列表，添加完整路径
            for file in files:
                if "path" in file and "size" in file:
                    file_list.append({
                        "path": file["path"],
                        "size": file["size"],
                        "mirror": mirror,
                        "type": file.get("type", "file")
                    })
            
            # 如果成功获取文件列表，跳出循环
            if file_list:
                break
                
        except Exception as e:
            print(f"从 {mirror} 获取文件列表失败: {e}")
    
    # 过滤出文件（不包括目录）
    return [f for f in file_list if f.get("type") == "file"]

def download_file_with_aria2(file_info, status_data):
    """使用aria2下载单个文件，支持断点续传"""
    file_path = file_info["path"]
    file_size = file_info.get("size", 0)
    mirror = file_info.get("mirror", MIRRORS[0])
    
    # 检查是否已下载成功
    if file_path in status_data["downloaded"]:
        local_path = os.path.join(SAVE_PATH, file_path)
        if os.path.exists(local_path) and os.path.getsize(local_path) == file_size:
            print(f"文件已下载完成，跳过: {file_path}")
            return True
    
    # 构建本地保存路径
    local_path = os.path.join(SAVE_PATH, file_path)
    local_dir = os.path.dirname(local_path)
    ensure_dir(local_dir)
    
    # 检查是否有部分下载的文件，这些文件会被断点续传
    has_partial_download = os.path.exists(f"{local_path}.aria2") or os.path.exists(local_path)
    if has_partial_download:
        print(f"发现部分下载文件，将进行断点续传: {file_path}")
    
    # 构建下载URL - HuggingFace文件下载链接
    encoded_path = quote(file_path)
    download_url = f"{mirror}/{MODEL_NAME}/resolve/main/{encoded_path}?download=true"
    
    # 构建aria2c命令行参数
    cmd = ['aria2c']
    for key, value in ARIA2_OPTS.items():
        cmd.append(f'--{key}={value}')
    
    # 添加输出路径
    cmd.append(f'--dir={local_dir}')
    cmd.append(f'--out={os.path.basename(local_path)}')
    
    # 添加下载URL
    cmd.append(download_url)
    
    try:
        print(f"开始下载: {file_path}")
        if has_partial_download:
            print(f"继续断点续传: {file_path}")
            
        # 使用subprocess.Popen获取实时输出
        process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            universal_newlines=True,
            bufsize=1
        )
        
        # 读取并打印输出
        for line in process.stdout:
            line = line.strip()
            if line and ("DL:" in line or "ERR" in line or "Download complete" in line):
                print(f"[{file_path}] {line}")
        
        # 等待进程结束
        process.wait()
        
        # 检查下载是否成功
        if process.returncode == 0:
            print(f"下载完成: {file_path}")
            # 验证文件大小
            if os.path.exists(local_path) and os.path.getsize(local_path) == file_size:
                # 更新下载状态
                if file_path not in status_data["downloaded"]:
                    status_data["downloaded"].append(file_path)
                if file_path in status_data["failed"]:
                    status_data["failed"].remove(file_path)
                save_download_status(status_data)
                return True
            else:
                print(f"文件大小验证失败: {file_path}")
                # 将此文件标记为失败
                if file_path not in status_data["failed"]:
                    status_data["failed"].append(file_path)
                save_download_status(status_data)
                return False
        else:
            print(f"下载失败: {file_path}")
            # 将此文件标记为失败
            if file_path not in status_data["failed"]:
                status_data["failed"].append(file_path)
            save_download_status(status_data)
            return False
    except Exception as e:
        print(f"下载过程出错: {e}")
        # 将此文件标记为失败
        if file_path not in status_data["failed"]:
            status_data["failed"].append(file_path)
        save_download_status(status_data)
        return False

def main():
    start_time = time.time()
    print(f"开始从HuggingFace下载 {MODEL_NAME} 模型到 {SAVE_PATH}")
    print(f"启用断点续传功能，支持中断后继续下载")
    
    # 确保下载目录存在
    ensure_dir(SAVE_PATH)
    
    # 加载下载状态
    status_data = load_download_status()
    print(f"已完成下载: {len(status_data['downloaded'])} 个文件，失败: {len(status_data['failed'])} 个文件")
    
    # 检查aria2是否已安装
    try:
        process = subprocess.run(['aria2c', '--version'], capture_output=True, check=True, text=True)
        version = process.stdout.split('\n')[0] if process.stdout else "未知版本"
        print(f"aria2c 已安装: {version}")
    except (subprocess.SubprocessError, FileNotFoundError):
        print("错误: aria2c 未安装。请先安装aria2: sudo apt install -y aria2")
        return
    
    # 获取文件列表
    all_files = get_file_list_from_huggingface()
    if not all_files:
        print("无法获取文件列表，下载终止")
        return
    
    # 优先处理失败的文件
    failed_files = [f for f in all_files if f["path"] in status_data["failed"]]
    remaining_files = [f for f in all_files if f["path"] not in status_data["downloaded"] and f["path"] not in status_data["failed"]]
    
    print(f"获取到 {len(all_files)} 个文件，其中:")
    print(f"- 已完成: {len(status_data['downloaded'])} 个")
    print(f"- 失败需重试: {len(failed_files)} 个")
    print(f"- 待下载: {len(remaining_files)} 个")
    
    # 合并失败文件和剩余文件，优先处理失败的
    files_to_download = failed_files + remaining_files
    
    if not files_to_download:
        print("所有文件已下载完成!")
        return
    
    print(f"开始使用aria2多线程下载，将处理 {len(files_to_download)} 个文件...")
    
    # 使用线程池执行下载任务
    with ThreadPoolExecutor(max_workers=MAX_PARALLEL_TASKS) as executor:
        # 提交所有下载任务
        future_to_file = {
            executor.submit(download_file_with_aria2, file_info, status_data): 
            file_info["path"] for file_info in files_to_download
        }
        
        # 收集结果
        success_count = 0
        fail_count = 0
        
        # 等待所有任务完成
        for future in future_to_file:
            file_path = future_to_file[future]
            try:
                if future.result():
                    success_count += 1
                else:
                    fail_count += 1
                    print(f"文件 {file_path} 下载失败")
            except Exception as e:
                fail_count += 1
                print(f"文件 {file_path} 处理异常: {e}")
    
    # 计算总耗时
    elapsed_time = time.time() - start_time
    print(f"下载完成! 成功: {success_count}, 失败: {fail_count}, 总耗时: {elapsed_time:.2f} 秒")
    
    # 最终状态汇总
    final_status = load_download_status()
    total_files = len(all_files)
    downloaded_files = len(final_status["downloaded"])
    completion_rate = (downloaded_files / total_files) * 100 if total_files > 0 else 0
    
    print(f"下载进度: {downloaded_files}/{total_files} ({completion_rate:.2f}%)")
    
    if fail_count == 0 and downloaded_files == total_files:
        print(f"所有文件下载成功! 模型保存在: {SAVE_PATH}")
    else:
        print(f"部分文件下载失败，可重新运行脚本继续断点续传")
        print(f"已下载文件将被跳过，只会下载失败或未下载的文件")

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n下载被用户中断，您可以随时重新运行脚本继续断点续传下载")
        sys.exit(1) 